from collections_abc import Indexable
from collections import Callable
import itertools
from numbers import Integral
from numpy.random import RandomState


Dataset = Indexable


def _check_index(length, index):
    return 0 <= index < length


class IgnoredIndexException(Exception):
    pass


class map(Indexable):
    def __init__(self, func, indexable, *more):
        assert isinstance(func, Callable)
        assert isinstance(indexable, Indexable)
        assert all(isinstance(i, Indexable) for i in more)

        self._func = func
        self._indexables = (indexable,) + more
        self._len = len(indexable)
        assert all(len(i) == self._len for i in more)

    def __len__(self):
        return self._len

    def __getitem__(self, index):
        assert _check_index(len(self), index)
        return self._func(*(i[index] for i in self._indexables))


class product(Indexable):
    def __init__(self, indexable_a, indexable_b, *more):
        assert isinstance(indexable_a, Indexable)
        assert isinstance(indexable_b, Indexable)
        assert all(isinstance(i, Indexable) for i in more)

        self._indexables = (indexable_a, indexable_b) + more
        self._table = tuple(itertools.product(*(range(len(i)) for i in self._indexables)))
        self._len = len(self._table)

    def __len__(self):
        return self._len

    def __getitem__(self, index):
        assert _check_index(len(self), index)
        return tuple(i[ti] for i, ti in zip(self._indexables, self._table[index]))


class concatenate(Indexable):
    def __init__(self, indexable_a, indexable_b, *more):
        assert isinstance(indexable_a, Indexable)
        assert isinstance(indexable_b, Indexable)
        assert all(isinstance(i, Indexable) for i in more)

        self._indexables = (indexable_a, indexable_b) + more
        self._out_table = tuple(itertools.chain.from_iterable(
            itertools.repeat(idx_i, len(i)) for idx_i, i in enumerate(self._indexables)
        ))
        self._in_table = tuple(itertools.chain.from_iterable(
            range(len(i)) for i in self._indexables
        ))
        self._len = len(self._out_table)

    def __len__(self):
        return self._len

    def __getitem__(self, index):
        assert _check_index(len(self), index)
        return self._indexables[self._out_table[index]][self._in_table[index]]


class choice(Indexable):
    def __init__(self, indexable, nb_samples, rng):
        assert isinstance(indexable, Indexable)
        assert isinstance(nb_samples, Integral) and nb_samples > 0
        assert isinstance(rng, RandomState)

        self._indexable = indexable
        self._nb_samples = nb_samples
        self._table = rng.choice(len(indexable), nb_samples, replace=False)
        self._len = nb_samples

    def __len__(self):
        return self._len

    def __getitem__(self, index):
        assert _check_index(len(self), index)
        return self._indexable[self._table[index]]


class repeat(Indexable):
    def __init__(self, indexable, times):
        assert isinstance(indexable, Indexable)
        assert isinstance(times, Integral) and times > 0

        self._indexable = indexable
        self._times = times
        self._len_of_indexable = len(indexable)
        self._len = self._len_of_indexable * times

    def __len__(self):
        return self._len

    def __getitem__(self, index):
        assert _check_index(len(self), index)
        return self._indexable[index % self._len_of_indexable]
